"""
-*- coding: utf-8 -*-
-*- test version: 3.7.4 -*-
@Time    : 2024/8/13 16:47
"""

import json
import asyncio
import websockets
from pydub import AudioSegment
import numpy as np

async def play_file(websocket):
    # Read the WAV file and send it in chunks of 10240 bytes
    audio = AudioSegment.from_file('10s.wav', format="wav").set_frame_rate(8000).set_sample_width(2).set_channels(1)

    chunk_size = 10240
    for i in range(0, len(audio.raw_data), chunk_size):
        chunk = audio.raw_data[i:i+chunk_size]
        await send_data(websocket, chunk)

def decode_s16_8000(data):
    # Assuming little-endian signed 16-bit integers
    audio_data = np.frombuffer(data, dtype=np.int16)
    return audio_data.tobytes()
    

async def send_data(websocket, data):
    await websocket.send(data)

async def handle_start(websocket, rpc):
    reply = {
        "jsonrpc": "2.0",
        "id": rpc["params"]["uuid"],
        # "id": rpc["id"],
        "result": {
            "code": 200,
            "message": "OK",
            "audio": "sendrecv"
        }
    }

    print("send", reply)
    await websocket.send(json.dumps(reply))
    await play_file(websocket)

async def handle_message(websocket, message):
    # print("handle_message")
    global packets

    if isinstance(message, bytes):  # Audio Data
        # print(message)
        # you can use 'ffmpeg -f s16le -ar 8000 -ac 1 -i test.raw test.wav' to convert test.raw into test.wav
        with open('test.raw', 'ab') as raw_file:
            raw_file.write(message)
    else:
        rpc = json.loads(message)
        print(rpc)

        if rpc["method"] == "start":
            await handle_start(websocket, rpc)


async def on_connection(websocket, path):
    try:
        async for message in websocket:
            await handle_message(websocket, message)
    finally:
        print("connection closed")

start_server = websockets.serve(on_connection, "0.0.0.0", 9001, ping_interval=None)

asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()
